# 复数CNN网络用于求解多参数估计问题
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class ComplexReLU(nn.Module):
    def forward(self, input):
        # 分别对实部和虚部应用ReLU
        real = F.relu(input.real)
        imag = F.relu(input.imag)
        return torch.complex(real, imag).to(dtype=torch.complex128).to(input.device)


class ComplexBatchNorm2d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
                 track_running_stats=True):
        super(ComplexBatchNorm2d, self).__init__()
        self.bn = nn.BatchNorm2d(num_features)

    def forward(self, X):
        # 将归一化后的实部和虚部合并成复数张量
        input_re_normalized = self.bn(torch.real(X).float())
        input_im_normalized = self.bn(torch.imag(X).float())
        output = torch.complex(input_re_normalized.to(dtype=torch.float64), input_im_normalized.to(dtype=torch.float64)).to(X.device)
        return output

class ComplexMaxPool2d(nn.MaxPool2d):
    def __init__(self, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False):
        super(ComplexMaxPool2d, self).__init__(kernel_size)
        mp_size, mp_stride = 4, 4
        self.mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def forward(self, X):
        # 将归一化后的实部和虚部合并成复数张量
        input_re_normalized = self.mp(torch.real(X).float())
        input_im_normalized = self.mp(torch.imag(X).float())
        output = torch.complex(input_re_normalized.to(dtype=torch.float64), input_im_normalized.to(dtype=torch.float64)).to(X.device)
        return output

class MpeCcnnModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 第1层：256*256*8 -> 64*64*16
        self.gen_l1()
        # 第2层：64*64*16 -> 16*16*32
        self.gen_l2()
        # 第3层：16*16*32 -> 4*4*64
        self.gen_l3()
        # 拉直
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(4 * 4 * 64 * 2, 256) # *2是为了将复数实部和虚部串接起来
        self.fc2 = nn.Linear(256, 12)

    def forward(self, x):
        a1 = self.l1_mp(self.l1_relu(self.l1_bn(self.l1_conv2d(x))))
        a2 = self.l2_mp(self.l2_relu(self.l2_bn(self.l2_conv2d(a1))))
        a3c = self.l3_mp(self.l3_relu(self.l3_bn(self.l3_conv2d(a2))))
        a3 = torch.hstack((torch.real(a3c).float(), torch.imag(a3c).float()))
        a4 = self.flatten(a3)
        a5 = F.relu(self.fc1(a4))
        a6 = F.relu(self.fc2(a5))
        return a6

    def gen_l1(self):
        in_channels = 8
        out_channels = 16
        kernel = (31, 31)
        stride = (1,1)
        dilation = 1
        self.l1_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation, dtype=torch.complex128)
        self.l1_bn = ComplexBatchNorm2d(out_channels) #nn.BatchNorm2d(out_channels)
        self.l1_relu = ComplexReLU() # nn.ReLU(True)
        mp_size, mp_stride = 4, 4
        self.l1_mp = ComplexMaxPool2d(mp_size, mp_stride, padding=0) # nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l2(self):
        in_channels = 16
        out_channels = 32
        kernel = (7, 7)
        stride = (1, 1)
        dilation = 1
        self.l2_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation, dtype=torch.complex128)
        self.l2_bn = ComplexBatchNorm2d(out_channels) # nn.BatchNorm2d(out_channels)
        self.l2_relu = ComplexReLU() # nn.ReLU(True)
        mp_size, mp_stride = 4, 4
        self.l2_mp = ComplexMaxPool2d(mp_size, mp_stride, padding=0) # nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l3(self):
        in_channels = 32
        out_channels = 64
        kernel = (3, 3)
        stride = (1, 1)
        dilation = 1
        self.l3_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation, dtype=torch.complex128)
        self.l3_bn = ComplexBatchNorm2d(out_channels) # nn.BatchNorm2d(out_channels)
        self.l3_relu = ComplexReLU() #nn.ReLU(True)
        mp_size, mp_stride = 4, 4
        self.l3_mp = ComplexMaxPool2d(mp_size, mp_stride, padding=0) # nn.MaxPool2d(mp_size, mp_stride, padding=0)